import logging

logger = logging.getLogger(__name__)

class Register:
    def __init__(self, config, user_unused_ops_list=None, lib=None):
        # lib is a instance of torch.library.Library
        # Some inference chips may not support the backward implementation of operators
        self.lib = lib

        # reg_key like 'CUDA', reg_bac_key like AutogradCUDA
        self.reg_key = "PrivateUse1"

        self.all_ops = []
        self.vendor_unused_ops_list = []
        self.unused_ops = user_unused_ops_list + self.vendor_unused_ops_list
        self.config = config
        self.config_filter()
        self.for_each()

    def config_filter(self):
        self.config = [
            item for item in self.config if item[1].__name__ not in self.unused_ops
        ]

    def register_impl(self, key, fn):
        device_key = self.reg_key
        self.all_ops.append(key)
        self.lib.impl(key, fn, device_key)

    def for_each(self):
        try:
            for key, func in self.config:
                self.register_impl(key, func)
        except Exception as e:
            raise RuntimeError(e, "An exception  occurred while replacing the customization operator.")

    def get_all_ops(self):
        return self.all_ops

    def get_unused_ops(self):
        return self.unused_ops

    def get_vendor_name(self):
        return "ascend"

    def get_current_device(self):
        return "npu"